import os
import math
import pickle
import numpy
from scipy import interpolate
from CalTorque import OpenCalTorqueXLS, GetSheetData
from scipy.optimize import minimize
from scipy.optimize import curve_fit


#____________________________________________________________________________

def GetRunData(path = 'data/'):
    datapath = path
    lsdir = os.listdir(datapath)

    for f in lsdir:
        if (f.split('.')[-1] not in ['xlsx', 'xls']):
            continue
    print f
    book = OpenCalTorqueXLS(datapath+f)
    data = GetSheetData(book)
    
    return data, f


def GetTestRunData(path):
    book = OpenCalTorqueXLS(path)
    data = GetSheetData(book)
    return data
    
    
def GetPklAverage(path = 'pickled_average_in.pkl'):
    file = open('pickled_average_in.pkl', 'rb')
    avgdata = pickle.load(file)
    file.close()

    #Xavg = numpy.array(avgdata['x'])    #len(Xavg) = 400
    #Yavg = numpy.array(avgdata['y'])
    
    return avgdata

def AverageTorqueInterp(x, offset = 0):            ### if put x as argument, would there be an offset arg too? OFFSET FOR AVERAGE
    x2 = []
    f = interpolate.interp1d(Xavg, Yavg)
    if offset != 0:
        for i, j in enumerate(x):
            x2.append(x[i] - offset)
    else:
        x2 = x

    #Xin = []
    #for i in x2:
    #    if i > 10 and i < 390:
    #        Xin.append(i)
    #print len(Xin)
    return f(x2)


def PolyTwoGauss(p, x):
    offset = 52            ### distance between peaks
    #p[0],p[1],p[2] = amplitude , center_x, sigma: for peak A
    #p[3],p[4] = amplitude, sigma: for peak B
    #p[5],p[6],p[7] = ax^2 + bx + c
    #p[8] = 
    x0 = x[0]
    y = numpy.zeros(len(x0))
    for i, j in enumerate(x0):  #index 1 is out of bounds for axis 0 with size 1
        y[i] = p[0]*math.exp(-0.5*((j-p[1])/p[2])**2.) + p[3]*math.exp(-0.5*((j-(p[1]+offset))/p[4])**2.) + p[6]*(j - p[5])**2 + p[7]*(j - p[5]) + p[8]
        if y[i] < 1e-5: 
            y[i] = 1e-5
    return y


def PolyFn(x, x0, a, b, c):     # assumes ydata = f(xdata, params)
   return a*(x-x0)**2 + b*(x-x0) + c

def Residual(p, y, x):
    modely = PolyTwoGauss(p, x)
    return numpy.sum((modely - y)**2)


### idk if binning the data is the best option - maybe just add something after it?
def BinData(x, y):
    bins = 200.0
    max_d = 400.0
    x_inputs = numpy.zeros(bins)
    xdata = numpy.linspace(0, max_d, bins)
    ydata = numpy.zeros(bins)

    for j, k in enumerate(xdata):
        pos_torque = []
        x0 = k - (max_d / bins) / 2.
        x1 = k + (max_d / bins) / 2.
        for i, p in enumerate(x):
           if p > x0 and p <= x1:
              pos_torque.append(y[i])
        #print len(pos_torque), k
        if len(pos_torque) > 0:
            ydata[j] += max(pos_torque)
            x_inputs[j] += 1.

    for i, n in enumerate(x_inputs):
        if n == 0:
            continue
        ydata[i] = ydata[i]/n
    return xdata, ydata


#delta - minimum difference between peak and following pts
def PeakFinder(y_ax, x_ax, lookahead = 5, delta = 1):
    max_peaks = {}
    max_peaks['maxpos'] = []
    max_peaks['max'] = []
    peaklocation = []
    max = y_ax[0]
    
    # Get peaks
    for i, y in enumerate(y_ax[: -lookahead]):
        if y > max:
            max = y
            maxpos = x_ax[i]  ### There has to be a better way of doing this.

        if y < max - delta and max != numpy.Inf:    # I think it might be a bad idea to layer this many ifs
            if y_ax[i:i + lookahead].max() < max:
                if maxpos not in max_peaks['maxpos']:   # for some reason the previous code returns multiple copies of the same peak - this prevents that
                    max_peaks['maxpos'].append(maxpos)
                    max_peaks['max'].append(max)
    
    # Filter peaks to only the ones in which we are interested
    for i, j in enumerate(max_peaks['maxpos']):
        if 300 <= j < 330:
            peaklocation.append(j)
        elif 350 < j <= 380:
            ### how to do an "only append the first peak" or just change 400 to 380?
            peaklocation.append(j)
    
    # maybe add something that raises error if there are more (or less) than 2 peaks?
    
    return peaklocation

def CorrectForOffset():

    Xavg = numpy.array(GetPklAverage()['x'])
    Yavg = numpy.array(GetPklAverage()['y'])
    x = Xavg[150:-5]
    y = Yavg[150:-5]
    
    poly_fit = curve_fit(PolyFn, x, y) #to get initial guess on the baseline
    poly_parms = poly_fit[0]

    # put in initial guesses now
    fit_parms = [5, 310, 5, 5, 5, poly_parms[0], poly_parms[1], poly_parms[2], poly_parms[3]]
    fit_bounds = [(0,20),(300,350), (1e-5, 40), (0, 20), (1e-5, 20), (None, None), (None,None), (None, None), (None, None)]

    result = minimize(Residual, fit_parms, args = (([y]), ([x])), method='TNC', bounds=fit_bounds,tol=0.0000005,options={'disp':True})
    new_parms = result.x

    print new_parms


#_______________________________________________________
#def PlotFit():

#PlotFit()
CorrectForOffset()
   
#Center x for average : 3.20185084e+02 = 320.185084
#Center of other : 372.185084
